package protocol

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
)

// 消息类型定义
const (
	MsgTypeHeartbeat  = 0x01 // 心跳消息
	MsgTypeAuth       = 0x02 // 认证消息
	MsgTypeRegister   = 0x03 // 服务注册
	MsgTypeUnregister = 0x04 // 服务注销
	MsgTypeData       = 0x05 // 数据传输
	MsgTypeConnect    = 0x06 // 连接请求
	MsgTypeDisconnect = 0x07 // 断开连接
	MsgTypeResponse   = 0x08 // 响应消息
	MsgTypeError      = 0x09 // 错误消息
)

// 消息头大小 (4字节长度 + 1字节类型)
const MessageHeaderSize = 5

// Message 通用消息结构
type Message struct {
	Type uint8       `json:"type"`
	Data interface{} `json:"data"`
}

// HeartbeatMessage 心跳消息
type HeartbeatMessage struct {
	Timestamp int64 `json:"timestamp"`
}

// AuthMessage 认证消息
type AuthMessage struct {
	Method string                 `json:"method"`
	Data   map[string]interface{} `json:"data"`
}

// RegisterMessage 服务注册消息
type RegisterMessage struct {
	ClientID string        `json:"client_id"`
	Services []ServiceInfo `json:"services"`
}

// ServiceInfo 服务信息
type ServiceInfo struct {
	Name       string `json:"name"`
	ServerPort int    `json:"server_port"`
	TargetHost string `json:"target_host"`
	TargetPort int    `json:"target_port"`
	Protocol   string `json:"protocol"`
}

// ConnectMessage 连接请求消息
type ConnectMessage struct {
	SessionID   string `json:"session_id"`
	ServiceName string `json:"service_name"`
	TargetHost  string `json:"target_host"`
	TargetPort  int    `json:"target_port"`
}

// DataMessage 数据传输消息
type DataMessage struct {
	SessionID string `json:"session_id"`
	Data      []byte `json:"data"`
}

// ResponseMessage 响应消息
type ResponseMessage struct {
	Success bool        `json:"success"`
	Message string      `json:"message"`
	Data    interface{} `json:"data,omitempty"`
}

// ErrorMessage 错误消息
type ErrorMessage struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// WriteMessage 发送消息到连接
func WriteMessage(conn net.Conn, msgType uint8, data interface{}) error {
	// 序列化消息数据
	var msgData []byte
	var err error

	if data != nil {
		msgData, err = json.Marshal(data)
		if err != nil {
			return err
		}
	}

	// 创建消息头：长度(4字节) + 类型(1字节)
	header := make([]byte, MessageHeaderSize)
	binary.BigEndian.PutUint32(header[0:4], uint32(len(msgData)))
	header[4] = msgType

	// 发送消息头
	if _, err := conn.Write(header); err != nil {
		return err
	}

	// 发送消息体
	if len(msgData) > 0 {
		if _, err := conn.Write(msgData); err != nil {
			return err
		}
	}

	return nil
}

// ReadMessage 从连接读取消息
func ReadMessage(conn net.Conn) (uint8, []byte, error) {
	// 读取消息头
	header := make([]byte, MessageHeaderSize)
	if _, err := io.ReadFull(conn, header); err != nil {
		return 0, nil, err
	}

	// 解析消息头
	length := binary.BigEndian.Uint32(header[0:4])
	msgType := header[4]

	// 读取消息体
	if length == 0 {
		return msgType, nil, nil
	}

	data := make([]byte, length)
	if _, err := io.ReadFull(conn, data); err != nil {
		return msgType, nil, err
	}

	return msgType, data, nil
}

// ParseMessage 解析消息数据
func ParseMessage(msgType uint8, data []byte, target interface{}) error {
	if len(data) == 0 {
		return nil
	}
	return json.Unmarshal(data, target)
}

// IsControlMessage 判断是否为控制消息
func (m *Message) IsControlMessage() bool {
	switch m.Type {
	case MsgTypeHeartbeat, MsgTypeAuth, MsgTypeRegister, MsgTypeUnregister, MsgTypeDisconnect, MsgTypeResponse:
		return true
	default:
		return false
	}
}

// IsDataMessage 判断是否为数据消息
func (m *Message) IsDataMessage() bool {
	return m.Type == MsgTypeData
}

// String 返回消息的字符串表示
func (m *Message) String() string {
	typeName := getMessageTypeName(m.Type)
	return fmt.Sprintf("Message{Type: %s(0x%02x), Length: %d}",
		typeName, m.Type, len(m.Data.([]byte)))
}

// getMessageTypeName 获取消息类型名称
func getMessageTypeName(msgType uint8) string {
	switch msgType {
	case MsgTypeHeartbeat:
		return "HEARTBEAT"
	case MsgTypeAuth:
		return "AUTH"
	case MsgTypeRegister:
		return "REGISTER"
	case MsgTypeUnregister:
		return "UNREGISTER"
	case MsgTypeData:
		return "DATA"
	case MsgTypeConnect:
		return "CONNECT"
	case MsgTypeDisconnect:
		return "DISCONNECT"
	case MsgTypeResponse:
		return "RESPONSE"
	case MsgTypeError:
		return "ERROR"
	default:
		return "UNKNOWN"
	}
}
